"""
专用脚本,用于构造本算法中指定的others类型数据集的数据输入格式

"""
import os
import shutil
import random
import xml.etree.ElementTree as ET



old_dir = "VOC2007"
new_dir = "flawdataset"
custom_classes = ("03040001", "03010001", "02020001", "02010001", "03020001", "03030001")
train_cls_dict = {}
for i, cls in enumerate(custom_classes):
    train_cls_dict[cls] = i


if os.path.exists(new_dir):
    shutil.rmtree(new_dir)

os.makedirs(os.path.join(new_dir, "annotations"))
os.makedirs(os.path.join(new_dir, "train2017"))
os.makedirs(os.path.join(new_dir, "test2017"))


def xml_reader(filename):
    """ Parse a PASCAL VOC xml file """
    tree = ET.parse(filename)
    size = tree.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    objects = []
    for obj in tree.findall('object'):
        obj_struct = {}
        obj_struct['name'] = obj.find('name').text
        bbox = obj.find('bndbox')
        obj_struct['bbox'] = [int(bbox.find('xmin').text),
                              int(bbox.find('ymin').text),
                              int(bbox.find('xmax').text),
                              int(bbox.find('ymax').text)]
        objects.append(obj_struct)
        assert obj_struct['bbox'][0] < obj_struct['bbox'][2]
        assert obj_struct['bbox'][1] < obj_struct['bbox'][3]
    return width, height, objects


def generate_records(obj, writable_path):
    str_data = writable_path
    for item in obj:
        xmin, ymin, xmax, ymax = item["bbox"]
        class_name = item["name"]
        # iscrowd是否重叠设为0
        iscrowd = 1
        spot = f" {xmin},{ymin},{xmax},{ymax},{train_cls_dict[class_name]},{iscrowd}"
        str_data += spot
    str_data += "\n"
    return str_data


def main():
    jpeg_dir = os.path.join(old_dir, "JPEGImages")
    xml_dir = os.path.join(old_dir, "Annotations")
    xml_list = os.listdir(xml_dir)
    random.shuffle(xml_list)

    trainval_percent = 0.8
    train_percent = 0.8

    ftrainval = open(os.path.join(new_dir, "annotations", "trainval.txt"), 'w')
    ftest = open(os.path.join(new_dir, "annotations", "test.txt"), 'w')
    ftrain = open(os.path.join(new_dir, "annotations", "train.txt"), 'w')
    fval = open(os.path.join(new_dir, "annotations", "val.txt"), 'w')

    num = len(xml_list)
    xml_index = range(num)
    tv = int(num * trainval_percent)
    tr = int(tv * train_percent)
    trainval = random.sample(xml_index, tv)
    train = random.sample(trainval, tr)

    for i in xml_index:
        xml_name = xml_list[i]

        w, h, obj = xml_reader(os.path.join(xml_dir, xml_name))

        pic_name = xml_list[i][:-4] + ".jpg"
        pic_path = os.path.join(jpeg_dir, pic_name)

        if i in trainval:
            # 构造记录的内容
            # train2017/0000001.jpg 0,259,401,459,7 35,28,324,201,2 0,30,59,80,2
            writable_path = os.path.join("train2017", pic_name)
            str_data = generate_records(obj, writable_path)
            ftrainval.write(str_data)
            if i in train:
                ftrain.write(str_data)
            else:
                fval.write(str_data)
            det_pic_path = os.path.join(new_dir, "train2017", pic_name)
            shutil.copy(pic_path, det_pic_path)
        else:
            writable_path = os.path.join("test2017", pic_name)
            str_data = generate_records(obj, writable_path)
            ftest.write(str_data)
            det_pic_path = os.path.join(new_dir, "test2017", pic_name)
            shutil.copy(pic_path, det_pic_path)


if __name__ == '__main__':
    main()
